#include "types.h"
#include "expandkey.h"
#include "h_func.h"
#include "matrix.h"
#include "field.h"
#include "tables.h"
#include "mul_tab.h"
#include "pht.h"
#include "rotate.h"

#include <assert.h>
#include <stdlib.h>
#include <string.h>

byte gf_256_14d_mul(byte lhs, byte rhs)
{
	return element_mul(lhs, rhs, gf_256_14d_mul_tab, gf_256_14d_element_num);
}


int pre_expand_key(byte* key, int key_size, key_type* des_key)
{
	static const int round = 16;
	static const word roconst = 0x01010101;
	int i;
	int ret_code = 0;
	assert(key_size >= 0);
	assert(key_size <= 32);

	int fixed_key_size ;
	if(key_size <= 16)
		fixed_key_size = 16;
	else if(key_size <= 24)
		fixed_key_size = 24;
	else
		fixed_key_size = 32;

	des_key->key_size = fixed_key_size;

	int k = fixed_key_size / 8;
	byte* padded_key = des_key->orginal_key;
	memcpy((void*)padded_key, (void*)key, key_size);
	memset( (void*)(padded_key + key_size), fixed_key_size - key_size, 0);

	word* me = (word*)malloc( k * sizeof(word) );
	word* mo = (word*)malloc( k * sizeof(word) );

	for( i = 0; i < k; ++i)
	{
		me[i] = join_word( padded_key + (i*2)*4 );
		mo[i] = join_word( padded_key + (i*2 + 1) * 4);
	}

	byte tmp[4];
	for( i = 0; i < k; ++i)
	{
		byte_matrix_mul(rs, 4, 8, padded_key + 8*i, 8, 1, tmp, 4, element_add,  gf_256_14d_mul);
		des_key->s_array[ k - i - 1] = join_word(tmp);
	}

	word tmpa,tmpb,tmpk1,tmpk2;
	for( i = 0; i < round + 4; ++i)
	{
		tmpa = h_func(2*i*roconst, me, k);
		tmpb = rol(h_func((2*i+1)*roconst, mo, k), 8);
		//printf("Key%d,A:%x,B:%x\n", i, tmpa, tmpb);
		pht(tmpa, tmpb, &tmpk1, &tmpk2);
		tmpk2 = rol(tmpk2, 9);
		
		des_key->expanded_key[2*i] = tmpk1;
		des_key->expanded_key[2*i+1] = tmpk2;
	//	printf("%.8x\t%.8x\n", tmpk1, tmpk2);

	}

	free(mo);
	free(me);
	return ret_code;
}

#include <stdio.h>
#include <stdlib.h>
#include <assert.h>

void test_pre_expand_key()
{
	printf("Start testing pre_expand_key.........\n");
	
	//------------------------------Testing 128bit key-----------------------------
	byte key[16] = {
		0
	};
	key_type des_key;

	word real_expanded_keys[] = {
	0x52c54dde        ,0x11f0626d,
	0x7cac9d4a        ,0x4d1b4aaa,
	0xb7b83a10        ,0x1e7d0beb,
	0xee9c341f        ,0xcfe14be4,
	0xf98ffef9        ,0x9c5b3c17,
	0x15a48310,        0x342a4d81,
	0x424d89fe,        0xc14724a7,
	0x311b834c,        0xfde87320,
	0x3302778f,        0x26cd67b4,
	0x7a6c6362,        0xc2baf60e,
	0x3411b994,        0xd972c87f,
	0x84adb1ea,        0xa7dee434,
	0x54d2960f,        0xa2f7caa8,
	0xa6b8ff8c,        0x8014c425,
	0x6a748d1c,        0xedbaf720,
	0x928ef78c,        0x0338ee13,
	0x9949d6be,        0xc8314176,
	0x07c07d68,        0xecae7ea7,
	0x1fe71844,        0x85c05c89,
	0xf298311e,        0x696ea672,
	};

	pre_expand_key(key, sizeof(key) / sizeof(key[0]), &des_key);

	int i = 0;
	for(i = 0; i < sizeof(real_expanded_keys) / sizeof(real_expanded_keys[0]); ++i)
	{
		assert(real_expanded_keys[i] == des_key.expanded_key[i]);
	}

	//--------------------------------Testing 192bits keys-----------------------------------------
	const char* char_key192 = "01 23 45 67 89 AB CD EF FE DC BA 98 76 54 32 10 00 11 22 33 44 55 66 77";
	const int key192_size = 192/8;
	byte key192[key192_size];
	for( i = 0; i < key192_size; ++i)
	{
		//bug sscanf of %x fill 4bytes but only give 1byte, this may corrupt statck 
		sscanf(char_key192+3*i, "%x", (key192 + i));
	}

	pre_expand_key(key192, key192_size, &des_key);
/*	for(i = 0; i < EXPANDED_KEY_SIZE / 8; ++i)
	{
		printf("%x %x\n", des_key.expanded_key[2*i], des_key.expanded_key[2*i+1]);
	}*/

	word key192_real_expanded_keys[] = {
		0x38394a24, 0xc36d1175,
		0xe802528f, 0x219bfeb4,
		0xb9141ab4, 0xbd3e70cd,
		0xaf609383, 0xfd36908a,
		0x3efb931 ,0x1d2ee7ec,
		0xa7489d55, 0x6e44b6e8,
		0x714ad667, 0x653ad51f,
		0xb6315b66, 0xb27c05af,
		0xa06c8140, 0x9853d419,
		0x4016e346, 0x8d1c0dd4,
		0xf05480be, 0xb6af816f,
		0x2d7dc789, 0x45b7bd3a,
		0x57f8a163, 0x2befda69,
		0x26ae7271, 0xc2900d79,
		0xed323794, 0x3d3ffd80,
		0x5de68e49, 0x9c3d2478,
		0xdf326fe3, 0x5911f70d,
		0xc229f13b, 0xb1364772,
		0x4235364d, 0xcec363a,
		0x57c8dd1f, 0x6a1ad61e,
	};
	for(i = 0; i < sizeof(key192_real_expanded_keys) / sizeof(key192_real_expanded_keys[0]); ++i)
	{
		assert(key192_real_expanded_keys[i] == des_key.expanded_key[i]);
	}


	//----------------------------------Testing 256bits keys-----------------------------------
	const char* char_key256 = "01 23 45 67 89 AB CD EF FE DC BA 98 76 54 32 10 00 11 22 33 44 55 66 77 88 99 AA BB CC DD EE FF";
	const int key256_size = 256/8;
	byte key256[key256_size];
	for( i = 0; i < key256_size; ++i)
	{
		sscanf(char_key256+3*i, "%x", key256 + i);
	}

	pre_expand_key(key256, key256_size, &des_key);
//	for(i = 0; i < EXPANDD_KEY_SIZE / 8; ++i)
//	{
//		printf("%x %x\n", des_key.expanded_key[2*i], des_key.expanded_key[2*i+1]);
//	}

	word key256_real_expanded_keys[] = {
	0x5ec769bf, 0x44d13c60,
	0x76cd39b1, 0x16750474,
	0x349c294b, 0xec21f6d6,
	0x4fbd10b4, 0x578da0ed,
	0xc3479695, 0x9b6958fb,
	0x6a7fbc4e, 0xbf1830b,
	0x61b5e0fb, 0xd78d9730,
	0x7c6cf0c4, 0x2f9109c8,
	0xe69ea8d1, 0xed99bdff,
	0x35dc0bbd, 0xa03e5018,
	0xfb18ea0b, 0x38bd43d3,
	0x76191781, 0x37a9a0d3,
	0x72427bea, 0x911cc0b8,
	0xf1689449, 0x71009ca9,
	0xb6363e89, 0x494d9855,
	0x590bbc63, 0xf95a28b5,
	0xfb72b4e1, 0x2a43505c,
	0xbfd34176, 0x5c133d12,
	0x3a9247f7, 0x9a3331dd,
	0xee7515e6, 0xf0d54dcd,
	};
	for(i = 0; i < sizeof(key192_real_expanded_keys) / sizeof(key192_real_expanded_keys[0]); ++i)
	{
		assert(key256_real_expanded_keys[i] == des_key.expanded_key[i]);
	}


	printf("Test pre_expand_key complete!\n");

}
